from PIL import Image
from torch.utils.data import Dataset
import glob, os
import torch
import xarray as xr
class MyDataset(Dataset):
    def __init__(self, number, images_folder_train_x, images_folder_train_y, transform=None):
        images = []  # 存储输入的12图片
        images2 = []

        lr_list = sorted(glob.glob(os.path.join(images_folder_train_x, '*')),
                         key=lambda x: int((os.path.basename(x).split('_')[4])))
        list = []
        for name in lr_list:
            name = name[10:]
            list.append(name)
        # listName = os.listdir(images_folder_train_x)
        print("train_x number", len(list))
        count = 0
        num = 0
        while num < number:

            if count == 12:
                count = 0
                num = num - 11
            else:
                images.append(list[num])
                num = num + 1
                count = count + 1

        lr_list2 = sorted(glob.glob(os.path.join(images_folder_train_y, '*')),
                          key=lambda x: int((os.path.basename(x).split('_')[4])))
        list2 = []
        for name in lr_list2:
            name = name[10:]
            list2.append(name)
        # listName2 = os.listdir(images_folder_train_y)
        print("train_y number", len(list2))
        count2 = 0
        num2 = 0
        while num2 < number:

            if count2 == 12:
                count2 = 0
                num2 = num2 - 11
            else:
                images2.append(list2[num2])
                num2 = num2 + 1
                count2 = count2 + 1
        self.images_folder_train_x = images_folder_train_x
        self.images_folder_train_y = images_folder_train_y
        self.transform = transform
        self.images = images
        self.images2 = images2
        # print("image[]的格式")
        print("images number", len(images))
        print("images2 number", len(images2))

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        filename = self.images[index]
        img = xr.open_dataset(os.path.join(self.images_folder_train_x, filename))
        if int(filename[26:32])<=198708:
            img = img['N07_ICECON']
            img = torch.from_numpy(img.values)
        elif int(filename[26:32])<=199112:
            img = img['F08_ICECON']
            img = torch.from_numpy(img.values)
        elif int(filename[26:32])<=199509:
            img = img['F11_ICECON']
            img = torch.from_numpy(img.values)
        elif int(filename[26:32])<=200712:
            img = img['F13_ICECON']
            img = torch.from_numpy(img.values)
        else:
            img = img['F17_ICECON']
            img = torch.from_numpy(img.values)



        filename2 = self.images2[index]
        img2 = xr.open_dataset(os.path.join(self.images_folder_train_y, filename2))
        if int(filename2[26:32])<=198708:
            img2 = img2['N07_ICECON']
            img2 = torch.from_numpy(img2.values)
        elif int(filename2[26:32])<=199112:
            img2 = img2['F08_ICECON']
            img2 = torch.from_numpy(img2.values)
        elif int(filename2[26:32])<=199509:
            img2 = img2['F11_ICECON']
            img2 = torch.from_numpy(img2.values)
        elif int(filename2[26:32])<=200712:
            img2 = img2['F13_ICECON']
            img2 = torch.from_numpy(img2.values)
        else:
            img2 = img2['F17_ICECON']
            img2 = torch.from_numpy(img2.values)

        return img, img2
